{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom agent\n",
    "This notebook goes through how to create your own custom agent.\n",
    "\n",
    "此笔记本介绍如何创建自己的自定义代理。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `get_word_length` with `{'word': 'eudca'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m5\u001b[0m\u001b[32;1m\u001b[1;3mThe word \"eudca\" has 5 letters.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'actions': [ToolAgentAction(tool='get_word_length', tool_input={'word': 'eudca'}, log=\"\\nInvoking: `get_word_length` with `{'word': 'eudca'}`\\n\\n\\n\", message_log=[AIMessageChunk(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_JQzVVsPwO3FdinYV0sf9XwDt', 'function': {'arguments': '{\"word\":\"eudca\"}', 'name': 'get_word_length'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls'}, id='run-1fba6d2c-0789-4078-8fe4-d4adf844b634', tool_calls=[{'name': 'get_word_length', 'args': {'word': 'eudca'}, 'id': 'call_JQzVVsPwO3FdinYV0sf9XwDt'}], tool_call_chunks=[{'name': 'get_word_length', 'args': '{\"word\":\"eudca\"}', 'id': 'call_JQzVVsPwO3FdinYV0sf9XwDt', 'index': 0}])], tool_call_id='call_JQzVVsPwO3FdinYV0sf9XwDt')],\n",
       "  'messages': [AIMessageChunk(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_JQzVVsPwO3FdinYV0sf9XwDt', 'function': {'arguments': '{\"word\":\"eudca\"}', 'name': 'get_word_length'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls'}, id='run-1fba6d2c-0789-4078-8fe4-d4adf844b634', tool_calls=[{'name': 'get_word_length', 'args': {'word': 'eudca'}, 'id': 'call_JQzVVsPwO3FdinYV0sf9XwDt'}], tool_call_chunks=[{'name': 'get_word_length', 'args': '{\"word\":\"eudca\"}', 'id': 'call_JQzVVsPwO3FdinYV0sf9XwDt', 'index': 0}])]},\n",
       " {'steps': [AgentStep(action=ToolAgentAction(tool='get_word_length', tool_input={'word': 'eudca'}, log=\"\\nInvoking: `get_word_length` with `{'word': 'eudca'}`\\n\\n\\n\", message_log=[AIMessageChunk(content='', additional_kwargs={'tool_calls': [{'index': 0, 'id': 'call_JQzVVsPwO3FdinYV0sf9XwDt', 'function': {'arguments': '{\"word\":\"eudca\"}', 'name': 'get_word_length'}, 'type': 'function'}]}, response_metadata={'finish_reason': 'tool_calls'}, id='run-1fba6d2c-0789-4078-8fe4-d4adf844b634', tool_calls=[{'name': 'get_word_length', 'args': {'word': 'eudca'}, 'id': 'call_JQzVVsPwO3FdinYV0sf9XwDt'}], tool_call_chunks=[{'name': 'get_word_length', 'args': '{\"word\":\"eudca\"}', 'id': 'call_JQzVVsPwO3FdinYV0sf9XwDt', 'index': 0}])], tool_call_id='call_JQzVVsPwO3FdinYV0sf9XwDt'), observation=5)],\n",
       "  'messages': [FunctionMessage(content='5', name='get_word_length')]},\n",
       " {'output': 'The word \"eudca\" has 5 letters.',\n",
       "  'messages': [AIMessage(content='The word \"eudca\" has 5 letters.')]}]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain.agents import tool\n",
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain.agents.format_scratchpad.openai_tools import format_to_openai_tool_messages\n",
    "from langchain.agents.output_parsers.openai_tools import OpenAIToolsAgentOutputParser\n",
    "from langchain.agents import AgentExecutor\n",
    "\n",
    "# Define the llm\n",
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo\", temperature=0)\n",
    "\n",
    "# Define the tool\n",
    "@tool\n",
    "def get_word_length(word: str) -> int:\n",
    "    \"\"\"Returns the length of a word.\"\"\"\n",
    "    return len(word)\n",
    "\n",
    "tools = [get_word_length]\n",
    "\n",
    "# define prompt template\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are very powerful assistant, but don't know current events\",\n",
    "        ),\n",
    "        (\"user\", \"{input}\"),\n",
    "        MessagesPlaceholder(variable_name=\"agent_scratchpad\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# How does the agent know what tools it can use?\n",
    "# bind tools to llm\n",
    "llm_with_tools = llm.bind_tools(tools)\n",
    "\n",
    "# define agent\n",
    "agent = (\n",
    "    {\n",
    "        \"input\": lambda x: x[\"input\"],\n",
    "        #  Convert (AgentAction, tool output) tuples into FunctionMessages        \n",
    "        #  formatting intermediate steps (agent action, tool output pairs) to input messages that can be sent to the model\n",
    "        \"agent_scratchpad\": lambda x: format_to_openai_tool_messages(  \n",
    "            x[\"intermediate_steps\"]\n",
    "        ),\n",
    "    }\n",
    "    | prompt\n",
    "    | llm_with_tools\n",
    "    | OpenAIToolsAgentOutputParser()  # converting the output message into an agent action/agent finish\n",
    ")\n",
    "\n",
    "# Create an agent executor by passing in the agent and tools\n",
    "agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)\n",
    "\n",
    "list(agent_executor.stream({\"input\": \"How many letters in the word eudca\"}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `get_word_length` with `{'word': 'eudca'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m5\u001b[0m\u001b[32;1m\u001b[1;3mThe word \"eudca\" has 5 letters.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'How many letters in the word eudca',\n",
       " 'output': 'The word \"eudca\" has 5 letters.'}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor.invoke({\"input\": \"How many letters in the word eudca\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='5', response_metadata={'token_usage': {'completion_tokens': 1, 'prompt_tokens': 15, 'total_tokens': 16}, 'model_name': 'gpt-3.5-turbo', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-6d52c824-76fc-4a8a-ba9c-052d9b3eab2f-0')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm.invoke(\"How many letters in the word educa\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adding memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\n",
      "Invoking: `get_word_length` with `{'word': 'educa'}`\n",
      "\n",
      "\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3m5\u001b[0m\u001b[32;1m\u001b[1;3mThe word \"educa\" has 5 letters.\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n",
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3m\"Educa\" is not a common English word. It seems to be a variation or abbreviation of the word \"education.\"\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'is that a real word?',\n",
       " 'chat_history': [HumanMessage(content='how many letters in the word educa?'),\n",
       "  AIMessage(content='The word \"educa\" has 5 letters.')],\n",
       " 'output': '\"Educa\" is not a common English word. It seems to be a variation or abbreviation of the word \"education.\"'}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.prompts import MessagesPlaceholder\n",
    "from langchain_core.messages import AIMessage, HumanMessage\n",
    "\n",
    "MEMORY_KEY = \"chat_history\"\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are very powerful assistant, but bad at calculating lengths of words.\",\n",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=MEMORY_KEY),\n",
    "        (\"user\", \"{input}\"),\n",
    "        MessagesPlaceholder(variable_name=\"agent_scratchpad\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "chat_history = []\n",
    "agent = (\n",
    "    {\n",
    "        \"input\": lambda x: x[\"input\"],\n",
    "        \"agent_scratchpad\": lambda x: format_to_openai_tool_messages(\n",
    "            x[\"intermediate_steps\"]\n",
    "        ),\n",
    "        \"chat_history\": lambda x: x[\"chat_history\"],\n",
    "    }\n",
    "    | prompt\n",
    "    | llm_with_tools\n",
    "    | OpenAIToolsAgentOutputParser()\n",
    ")\n",
    "agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)\n",
    "\n",
    "# first time\n",
    "input1 = \"how many letters in the word educa?\"\n",
    "result = agent_executor.invoke({\"input\": input1, \"chat_history\": chat_history})\n",
    "\n",
    "# second time\n",
    "chat_history.extend(\n",
    "    [\n",
    "        HumanMessage(content=input1),\n",
    "        AIMessage(content=result[\"output\"]),\n",
    "    ]\n",
    ")\n",
    "agent_executor.invoke({\"input\": \"is that a real word?\", \"chat_history\": chat_history})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain0_1",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
